import matplotlib.pyplot as plt
import torchvision.transforms
from torch.utils.data import Dataset
import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from torchvision import datasets
from torch.utils.data import DataLoader


class MyData(Dataset):

    def __init__(self,data_path):
        self.datas = np.loadtxt(data_path,dtype=str,skiprows=1)

    def __getitem__(self, idx):
        data = [float(i) for i in self.datas[idx].split(',')]
        if len(data) == 785:
            label = int(data[0])
            simple = torch.Tensor(data[1:])
            return simple,label
        elif len(data) == 784:
            sample = torch.Tensor(data)
            return sample

    def __len__(self):
        return len(self.datas)

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1,1,3),
            nn.Flatten(),
            nn.Linear(26*26,1024),
            nn.ReLU(),
            nn.Linear(1024,1024),
            nn.ReLU(),
            nn.Linear(1024, 512),
            nn.ReLU(),
            nn.Linear(512,10)
        )

    def forward(self,x):
        return self.net(x)
if __name__ == '__main__':
    mnist_data = datasets.MNIST('./dataset', True, download=False, transform=torchvision.transforms.ToTensor())
    mnist_data_test = datasets.MNIST('./dataset', False, download=False, transform=torchvision.transforms.ToTensor())
    dataloader = DataLoader(mnist_data, batch_size=256,shuffle=True)
    dataloader_test = DataLoader(mnist_data_test, batch_size=64,shuffle=True)
    model = Model()
    model.load_state_dict(torch.load('./model_para_2.pth'))
    model.cuda()
    loss_fn = nn.CrossEntropyLoss()
    loss_fn.cuda()
    optimizer = torch.optim.SGD(model.parameters(),lr=0.05)

    epoch = 10
    # losses = []
    for i in range(epoch):
        print('正在训练第{}代'.format(i+1))
        loss_sum = 0
        current_num = 0
        for imgs,target in dataloader:
            imgs = imgs.cuda()
            outputs = model(imgs)
            target = target.cuda()
            current_num += (target == torch.argmax(outputs,dim=1)).sum().item()
            loss = loss_fn(outputs,target)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            loss_sum += loss.item()
        loss_ = loss_sum/len(mnist_data)
        print('第{}代损失为 {}'.format(i+1,loss_))
        print('正确率为{}'.format(current_num/len(mnist_data)))

        with torch.no_grad():
            current_num = 0
            loss_sum = 0
            for imgs,target in dataloader_test:
                imgs = imgs.cuda()
                outputs = model(imgs)
                target = target.cuda()
                current_num += (target == torch.argmax(outputs,dim=1)).sum().item()
                loss = loss_fn(outputs,target)
                # optimizer.zero_grad()
                # loss.backward()
                # optimizer.step()
                loss_sum += loss.item()
            loss_ = loss_sum/len(mnist_data_test)
            print('第{}代测试损失为 {}'.format(i+1,loss_))
            print('正确率为{}'.format(current_num/len(mnist_data_test)))

        # losses.append(loss_)
    # plt.plot(losses)
    # plt.show()
    torch.save(model.state_dict(),'model_para_2.pth')
    '''
    第50代损失为 1.4838810846495476
    第10代损失为 1.5223927792082441
    '''